import { Tensor } from 'onnxruntime-web'
import ndarray from 'ndarray'
import ops from 'ndarray-ops'
import _ from 'lodash'
import { softmax } from '@/utils/math'


// transform image to tensor
export function transformImage(imageData, width, height){
  // convert to tensor
  //创建输入图像的ndarray对象(dataTensor)，它是一个三维数组，表示图像的宽度、高度和通道数。
  const dataTensor = ndarray(new Float32Array(imageData.data), [width, height, 4])
  //创建输出图像的ndarray对象(dataProcessedTensor)，它是一个四维数组，表示一批处理后的图像，维度为[1, 3, 图像宽度, 图像高度]。
  const dataProcessedTensor = ndarray(new Float32Array(width * height * 3), [1, 3, width, height])

  // 归一化处理。将输入图像的通道数据复制到输出图像的对应位置，并进行归一化操作，实现图像的预处理。
  ops.assign(dataProcessedTensor.pick(0, 0, null, null), dataTensor.pick(null, null, 0)) // 将R通道数据复制到处理后的张量中的对应位置
  ops.assign(dataProcessedTensor.pick(0, 1, null, null), dataTensor.pick(null, null, 1)) // 将G通道数据复制到处理后的张量中的对应位置
  ops.assign(dataProcessedTensor.pick(0, 2, null, null), dataTensor.pick(null, null, 2)) // 将B通道数据复制到处理后的张量中的对应位置

  // 数据归一化到[0, 1]范围
  ops.divseq(dataProcessedTensor, 224)
  // 减去RGB通道的均值
  ops.subseq(dataProcessedTensor.pick(0, 0, null, null), 0.485)
  ops.subseq(dataProcessedTensor.pick(0, 1, null, null), 0.456)
  ops.subseq(dataProcessedTensor.pick(0, 2, null, null), 0.406)
  // 数据归一化到均值为0，标准差为1的范围
  ops.divseq(dataProcessedTensor.pick(0, 0, null, null), 0.229)
  ops.divseq(dataProcessedTensor.pick(0, 1, null, null), 0.224)
  ops.divseq(dataProcessedTensor.pick(0, 2, null, null), 0.225)

  // 创建一个新的张量对象并复制处理后的数据
 /* const tensor = new Tensor('float32', new Float32Array(width * height * 3), [1, 3, width, height])
  ;(tensor.data).set(dataProcessedTensor.data)*/
  console.log(dataProcessedTensor)
  const inputTensor = new Tensor('float32',new Float32Array(dataProcessedTensor.data), [1, 3, width,height]);
  console.log(inputTensor)

  return inputTensor;
}

// preprocess image将图像从URL转换为ImgData
export async function preprocess  (inputImageUrl) {
  const width = 224
  const height = 224
  // 从URL加载图像
  const image = await new Promise((resolve, reject) => {
    const image = new Image()
    image.onload = () => resolve(image)
    image.onerror = reject
    image.src = inputImageUrl
  })
  console.log(inputImageUrl)
  // 创建一个画布用于调整图像大小
  const canvas = document.createElement('canvas')
  canvas.width = width
  canvas.height = height
  const ctx = canvas.getContext('2d')
  // 在画布上绘制调整大小后的图像
  ctx.drawImage(image, 0, 0, width, height)
  // 获取绘制后的图像数据
  const imageData = ctx.getImageData(0, 0, width, height)
  // 如果成功获取图像数据，则将其转换为ImgData并进行处理

  if (imageData) {
    console.log(imageData)
    console.log(transformImage(imageData, width, height))
    return transformImage(imageData, width, height)
  } else {
    return null
  }

}
export async function RecordPreprocess  (fileName) {
  const width = 224
  const height = 224
  // 从URL加载图像
  console.log(fileName)

  const image = await new Promise((resolve, reject) => {
    const image = new Image()
    image.setAttribute('crossOrigin', 'anonymous')
    image.onload = () => resolve(image)
    image.onerror = reject
    image.src ='http://localhost:9090/files/'+fileName
  })

  console.log(image)
  // 创建一个画布用于调整图像大小
  const canvas = document.createElement('canvas')
  canvas.width = width
  canvas.height = height
  const ctx = canvas.getContext('2d')
  // 在画布上绘制调整大小后的图像
  ctx.drawImage(image, 0, 0, width, height)
  // 获取绘制后的图像数据
  const imageData = ctx.getImageData(0, 0, width, height)
  console.log(imageData);
  // 如果成功获取图像数据，则将其转换为ImgData并进行处理

  if (imageData) {
    console.log(imageData)
    console.log(transformImage(imageData, width, height))
    return transformImage(imageData, width, height)
  } else {
    return null
  }

}

// 后处理：将模型输出映射为类别及其概率
export function postprocess (
    output,                  // 模型的输出张量
    classIndices, // 类别索引映射
    k = 5                            // 返回前k个最高概率结果，默认为5
)  {
  const outputData = output.data
  if (!outputData || outputData.length === 0) {
    return []  // 如果输出数据为空，直接返回空数组
  }

  // 对输出数据应用 softmax 函数,输出为概率
  const outputSoftmax = softmax(Array.prototype.slice.call(outputData))
  const probs = _.isTypedArray(outputSoftmax)
      ? Array.prototype.slice.call(outputSoftmax)
      : outputSoftmax

  // 对概率进行排序，得到最高概率的前 k 个结果
  const sorted = _.reverse(
      _.sortBy(
          probs.map((prob, index) => [prob, index]),
          (probIndex) => probIndex[0]
      )
  )

  // 获取前 k 个最高概率的类别及其概率值
  const topK = _.take(sorted, k).map((probIndex) => {
    return {
      class: classIndices[probIndex[1]] ?? '', // 使用类别索引映射获取类别名称
      probability: probIndex[0]                // 概率值
    }
  })

  return topK  // 返回前 k 个最高概率结果
}


// 模型预测函数
export async function  predict (
    model, // ONNX Runtime InferenceSession 对象，用于模型推理
    inputTensor ,             // 输入图像的张量表示
    classIndices // 类别索引映射
){
  console.log("done")
  if (model && inputTensor && classIndices) {
    const startTime = new Date().getTime() // 记录推理开始时间
    const output = await model.run({ input: inputTensor }) // 进行模型推理，获取输出结果
    const results = postprocess(output.output, classIndices, 5) // 对输出结果进行后处理，获取前5个最高概率的类别及其概率

    // 如果存在处理后的结果
    if (results.length > 0) {
      return {
        data: results.map((result) => result.probability), // 返回结果中的概率值数组
        labels: results.map((result) => result.class),      // 返回结果中的类别名称数组
        timeCost: new Date().getTime() - startTime          // 计算推理所耗费的时间
      }
    }
    return null // 如果没有结果，返回空
  }
  return null // 如果输入参数为空，返回空
}

